import os
import redis
import time
import json
from multiprocessing import Pool
import sys
import subprocess
import datetime
sys.path.append("../")
from logs.log import setup_custom_logger

from DSSM.run_update_user_item_class import RecEmbeddingUpdater  # 从你的文件中导入类
from utils.redis_utils import _connect_redis
# Redis 连接配置
redis_client = _connect_redis()

# 显卡数量（根据你的实际硬件配置调整）
NUM_GPUS = 4  # 假设有 4 张显卡，可根据实际情况修改

# 所有 group_name 列表（根据你的需求调整）
GROUP_NAMES = ["JXSTNU", "SUDA"]  # 示例列表，替换为实际的 group_name

# 数据库配置文件路径
DB_CONFIG_PATH = "../global_configs/mysql_db_connect_config.json"

# 任务所需的最小显存（单位：MB）
MEMORY_THRESHOLD = 1024 * 44  # 例如 44GB，可根据任务需求调整


# 日志系统初始化
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
LOG_DIR = '../Logfiles'
file_name = str(datetime.datetime.now().strftime('%Y-%m-%d'))
my_logger = setup_custom_logger(os.path.join(LOG_DIR, '%s.log' % file_name), log_level="INFO")

def get_gpu_memory_usage():
    """
    获取每张显卡的显存使用情况。
    返回一个列表，列表中的每个元素是对应显卡的已用显存（MB）。
    """
    try:
        output = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,nounits,noheader"]
        ).decode('utf-8').strip().split('\n')
        memory_used = [int(x) for x in output]
        return memory_used
    except Exception as e:
        my_logger.info(f"获取显存使用情况失败: {e}")
        return [0] * NUM_GPUS  # 失败时返回全0，避免程序崩溃

def get_available_gpu():
    """
    获取一张可用的显卡 ID。
    条件：显卡未被占用且可用显存大于 MEMORY_THRESHOLD。
    如果没有可用显卡，则等待直到有显卡满足条件。
    """
    while True:
        memory_used = get_gpu_memory_usage()
        for gpu_id in range(NUM_GPUS):
            # 检查显卡是否在 "used_gpus" 集合中
            if not redis_client.sismember("used_gpus", gpu_id):
                # 获取显卡的总显存
                total_memory = int(subprocess.check_output(
                    ["nvidia-smi", "--query-gpu=memory.total", "--format=csv,nounits,noheader", "-i", str(gpu_id)]
                ).decode('utf-8').strip())
                available_memory = total_memory - memory_used[gpu_id]
                if available_memory > MEMORY_THRESHOLD:
                    # 占用该显卡
                    redis_client.sadd("used_gpus", gpu_id)
                    return gpu_id
        # 如果没有可用显卡，等待 1 秒后重试
        time.sleep(1)

def release_gpu(gpu_id):
    """
    释放指定的显卡资源。
    """
    redis_client.srem("used_gpus", gpu_id)

def worker(group_name):
    """
    工作函数：处理一个 group_name 的任务。
    - 获取可用显卡
    - 设置 CUDA_VISIBLE_DEVICES 环境变量
    - 执行 RecEmbeddingUpdater 的更新操作
    - 完成后释放显卡
    """
    # 获取可用显卡
    gpu_id = get_available_gpu()
    try:
        # 设置环境变量以指定使用的显卡
        os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
        
        # 加载数据库配置
        with open(DB_CONFIG_PATH, 'r') as file:
            db_config = json.load(file)
        
        # 实例化 RecEmbeddingUpdater 并执行更新
        updater = RecEmbeddingUpdater(
            db_config=db_config,
            group_name=group_name,
            milvus_host="10.240.1.3",
            milvus_port="19530"
        )
        result = updater.update_embeddings_for_group()
        my_logger.info(f"任务 {group_name} 完成，结果: {result}")
    except Exception as e:
        my_logger.error(f"任务 {group_name} 执行失败: {e}")
    finally:
        # 确保释放显卡资源
        release_gpu(gpu_id)

def main():
    """
    主函数：初始化任务队列并启动进程池。
    """
    # 将所有 group_name 加入 Redis 队列
    for group in GROUP_NAMES:
        redis_client.lpush("task_queue", group)
        my_logger.info(f"任务 {group} 已加入队列")

    # 初始化 used_gpus 集合（清空之前的状态）
    redis_client.delete("used_gpus")

    # 创建进程池，进程数等于显卡数量
    with Pool(processes=NUM_GPUS) as pool:
        while True:
            # 从 Redis 队列中阻塞式获取任务
            try:
                result = redis_client.blpop("task_queue", timeout=10)  # 10秒超时
                if result is None:
                    my_logger.warning("队列为空，退出调度")  # 这里如果你的任务场景是持续运行的系统，任务可能会在未来某个时间点动态加入队列，那么你可能希望程序继续等待，而不是退出。这样可以确保新任务到来时能够被及时处理。
                    break
                else:
                    _, group_name = result  # blpop 返回 (队列名, 值) 的元组
                    group_name = group_name.decode('utf-8')
                    my_logger.info(f"从队列中获取任务: {group_name}")
                    # 异步执行任务
                    pool.apply_async(worker, (group_name,))
            except Exception as e:
                my_logger.error(f"获取任务时出错: {e}")
                break
        
        # 等待所有任务完成
        pool.close()
        pool.join()

if __name__ == "__main__":
    main()